/*
   Copyright (c) 2011, The Chinese University of Hong Kong

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

/*
   This software contains source code provided by NVIDIA Corporation.
*/

/**
 * @file agrep.cpp
 *
 * @brief This is a traditional implementation of agrep algorithm,
 * which is used for approximate string matching.
 *
 * @author Hongjian Li, The Chinese University of Hong Kong.
 *
 * Copyright (C) 2011 Hongjian Li, The Chinese University of Hong Kong.
 */
#include <cutil_inline.h>	/** Used only for timing. */

// About corpus.
#define MAX_CORPUS_FILE_COUNT 40	/**< Maximum number of corpus files, i.e. maximum number of lines in Corpus.txt. */
#define MAX_CORPUS_FILE_PATH_LENGTH 100	/**< Maximum length of corpus file path, i.e. maximum length of lines in Corpus.txt. */
char corpus_file[MAX_CORPUS_FILE_COUNT][MAX_CORPUS_FILE_PATH_LENGTH + 1];	/**< Corpus filenames. */
unsigned int corpus_file_count;		/**< Actual number of corpus files. */

// About query. One query consists of a pattern followed by an edit distance.
#define MAX_PATTERN_COUNT 10000	/**< Maximum number of patterns, i.e. up to this many lines of matches will be stored into Result.csv. */
#define MAX_PATTERN_LENGTH 64	/**< Maximum pattern length. A 64-bit machine word has 64 bits. */
char pattern[MAX_PATTERN_COUNT][MAX_PATTERN_LENGTH + 2];	/**< Patterns. */
unsigned int edit_distance[MAX_PATTERN_COUNT];	/**< Number of primitive operations necessary to convert one string into an exact copy of the other. Primitive operations include insertion, deletion, and substitution of one character. */
unsigned int query_count;	/**< Actual number of queries. */

// About sequence.
#define MAX_SEQUENCE_COUNT 10000/**< Maximum number of sequences in Corpus.txt. Note that one .fa file might contain more than one sequence if it is not an assembled chromosome. */
#define MAX_LINE_LENGTH 1000	/**< Maximum length of lines of fasta files. Fasta files are scanned line by line in this program. This value should be larger than or at least equal to the length of fasta header, which far exceeds 70 in most cases. */
char sequence_header[MAX_SEQUENCE_COUNT][MAX_LINE_LENGTH];	/**< Headers of sequences. */
unsigned int sequence_length[MAX_SEQUENCE_COUNT];	/**< Lengthes of sequences. */
unsigned int sequence_cumulative_length[MAX_SEQUENCE_COUNT + 1];	/**< Cumulative lengths of sequences, i.e. 1) sequence_cumulative_length[0] = 0; 2) sequence_cumulative_length[sequence_index + 1] = sequence_cumulative_length[sequence_index] + sequence_length[sequence_index]; */
unsigned int sequence_count;	/**< Actual number of sequences. */

// About nucleotide.
#define CHARACTER_CARDINALITY 4	/**< One nucleotide is either A, C, G, or T. */
#define MAX_SCODON_COUNT 220000000	/**< Maximum number of special codons. The current value is 0.22G, therefore the special codon array can store up to 3.52G (=0.22G*16) base pairs. */
unsigned int scodon[MAX_SCODON_COUNT];	/**< The entire genomic nucleotides are stored into this array, one element of which, i.e. one 32-bit unsigned int, can store up to 16 nucleotides because one nucleotide can be uniquely represented by two bits since it must be either A, C, G, or T. One unsigned int is called a special codon, or scodon for short, because it is similar to codon, in which three consecutive characters of mRNA determine one amino acid of resulting protein. */
unsigned int character_count;	/**< Actual number of characters. */ 
unsigned int scodon_count;	/**< Actual number of special codons. scodon_count = (character_count + 16 - 1) >> 4; */

// About result.
#define MAX_MATCH_COUNT 10000	/**< Maximum number of matches of one single query. */

// About implementation.
#define K 9	/**< Maximum edit distance supported. Used to declare matching tables r[K + 1]. */
#define MAX_UNSIGNED_INT	0xffffffffUL	/**< The maximum value of an unsigned int. */
#define MAX_UNSIGNED_LONG_LONG	0xffffffffffffffffULL	/**< The maximum value of an unsigned long long. */

/**
 * Encode a character to its 2-bit binary representation.
 * The last two but one bits are different for A, C, G, and T respectively.
 * Note that some genomes contain 'N', which will be treated as 'G' in this encoding function.
 *
 * 'A' = 65 = 01000<b>00</b>1
 *
 * 'C' = 67 = 01000<b>01</b>1
 *
 * 'G' = 71 = 01000<b>11</b>1
 *
 * 'N' = 78 = 01001<b>11</b>0
 *
 * 'T' = 84 = 01010<b>10</b>0
 * @param[in] character The character to be encoded.
 * @return The 2-bit binary representation of given character.
 */
inline unsigned int encode(char character)
{
	return (character >> 1) & 3;
}

/**
 * Parse the corpus container file. It is the first argument to the program. By default it is Corpus.txt.
 * @param[in] corpus_container_file The container file specifying all the corpus files.
 */
void parseCorpusContainerFile(const char *corpus_container_file)
{
	FILE *corpus_container_file_handle = fopen(corpus_container_file, "r");	// The handle to corpus container file.
	corpus_file_count = 0;
	while (fgets(corpus_file[corpus_file_count], MAX_CORPUS_FILE_PATH_LENGTH, corpus_container_file_handle) != NULL)
	{
		corpus_file[corpus_file_count][strlen(corpus_file[corpus_file_count])-1] = '\0';
		corpus_file_count++;
	}
	fclose(corpus_container_file_handle);
}

/**
 * Build the special codon array. Corpus files are enumerated. Statistics of the corpus are calculated.
 */
void buildSpecialCodonArray()
{
	FILE *corpus_file_handle;	// The handle to corpus file.
	unsigned int corpus_file_index; // Used to enumerate corpus_file[].
	char line_buffer[MAX_LINE_LENGTH];	// fgets(line_buffer, MAX_LINE_LENGTH, corpus_file_handle);
	unsigned int line_buffer_length;	// Length of line_buffer.
	unsigned int line_buffer_index;	// Used to enumerate line_buffer[].
	int sequence_index;	// Used to enumerate sequence_header[] and sequence_length[].
	unsigned int character_index;	// Index of the current character across all the sequences of the entire corpus.
	unsigned int character_index_lowest_four_bits;	// The lowest four bits of character_index;
	unsigned int scodon_buffer;	// 16 consecutive characters will be accommodated into one 32-bit unsigned int.

	sequence_cumulative_length[0] = 0;
	sequence_index = -1;
	character_index = 0;
	scodon_buffer = 0;
	for (corpus_file_index = 0; corpus_file_index < corpus_file_count; corpus_file_index++)
	{
		corpus_file_handle = fopen(corpus_file[corpus_file_index], "r");
		while (fgets(line_buffer, MAX_LINE_LENGTH, corpus_file_handle) != NULL)
		{
			line_buffer_length = strlen(line_buffer) - 1; // Exclude the last character which is '\n' of ASCII code '\10'.
			if (line_buffer[0] == '>') // Header line
			{
				if (++sequence_index > 0) // Not the first sequence.
				{
					sequence_cumulative_length[sequence_index] = character_index;
					sequence_length[sequence_index - 1] = character_index - sequence_cumulative_length[sequence_index - 1];
				}
				memcpy(sequence_header[sequence_index], line_buffer, line_buffer_length);
				sequence_header[sequence_index][line_buffer_length] = '\0';
			}
			else
			{
				// Use line_buffer_index to enumerate line_buffer, equivalent to foreach (char character in line_buffer)
				for (line_buffer_index = 0; line_buffer_index < line_buffer_length; line_buffer_index++)
				{
					character_index_lowest_four_bits = character_index & 15;
					scodon_buffer |= encode(line_buffer[line_buffer_index]) << (character_index_lowest_four_bits << 1); // Earlier characters reside in lower bits, while later characters reside in higher bits.
					if (character_index_lowest_four_bits == 15) // The buffer is full. Flush it to the special codon array. Note that there is no need to clear scodon_buffer.
					{
						scodon[character_index >> 4] = scodon_buffer;
						scodon_buffer = 0;
					}
					character_index++;
				}
			}
		}
		fclose(corpus_file_handle);
	}
	// Calculate statistics.
	sequence_count = sequence_index + 1;
	character_count = character_index;
	sequence_cumulative_length[sequence_count] = character_count;
	sequence_length[sequence_index] = character_count - sequence_cumulative_length[sequence_index];
	scodon_count = (character_count + 16 - 1) >> 4;
	scodon[scodon_count - 1] = scodon_buffer;	// There are some corpus nucleotides in the special codon buffer, flush it. Now the last special codon might have zeros in its most significant bits. Don't treat such zeros as 'A's.
}

/**
 * Parse the query file. It is the second argument to the program. By default it is Query.txt.
 * Each line of the query file represents a query, which consists of a pattern followed by an edit distance.
 * @param[in] query_file The query file specifying all the queries.
 */
void parseQueryFile(const char *query_file)	// Save the queries into pattern[] and edit_distance[]. Save query_count.
{
	FILE *query_file_handle = fopen(query_file, "r");	// The handle to query file.
	unsigned int query_length;	// The length of each query.
	query_count = 0;
	while (fgets(pattern[query_count], MAX_PATTERN_LENGTH + 3, query_file_handle) != NULL)
	{
		query_length = strlen(pattern[query_count]) - 1; // Exclude the last character which is '\n' of ASCII code '\10'.
		edit_distance[query_count] = pattern[query_count][query_length - 1] - 48;	// char '0' is '\48'.
		pattern[query_count][query_length - 1] = '\0';
		query_count++;
	}
	fclose(query_file_handle);
}

/**
 * The entry of the agrep program.
 * @param[in] argc Number of arguments to the program. It should be equal to 5. Command: agrep Corpus.txt Query.txt Result.csv Log.csv
 * @param[in] argv Arguments to the program.
 */
int main(int argc, char** argv)
{
	// About query.
	const char *current_pattern;	// The pattern of current query.
	unsigned int m;	// Pattern length.
	unsigned int k;	// Edit distance.

	// About agrep algorithm.
	unsigned int       mask_array_32[CHARACTER_CARDINALITY];	// The 32-bit mask array of pattern.
	unsigned long long mask_array_64[CHARACTER_CARDINALITY];	// The 64-bit mask array of pattern.
	unsigned int       test_bit_32;	// The test bit for determining matches of patterns of length 32.
	unsigned long long test_bit_64;	// The test bit for determining matches of patterns of length 64.

	// About result.
	unsigned int match_sequence [MAX_MATCH_COUNT];	// The sequence  indexes of matches.
	unsigned int match_character[MAX_MATCH_COUNT];	// The character indexes of matches, i.e. ending positions of matches.
	unsigned int match_count;	// Actual number of matches.
	FILE *result_file_handle;	// The handle to result file "Result.csv".
	FILE *log_file_handle;		// The handle to log file "Log.csv".

	// About timing.
	unsigned int timer;	// Each timer has a unique id.
	float query_time;	// The querying time of one single query.

	printf("agrep v1.3 (Build January 13, 2011)\n");
	printf("Hongjian Li (hjli@cse.cuhk.edu.hk)\n");
	if (argc == 1) // argc should be equal to 5 in order to collect all the required arguments.
	{
		printf("Usage: agrep [InputCorpus] [InputQuery] [OutputResult] [OutputLog]\n");
		printf("e.g.:  agrep Corpus.txt Query.txt Result.csv Log.csv\n");
		return 0;
	}

	cutilCheckError(cutCreateTimer(&timer));	// Create a timer.
	printf("Reading and parsing corpus files and queries ... ");
	cutilCheckError(cutStartTimer(timer));	// Start timing.
	parseCorpusContainerFile(argv[1]);	// Parse the corpus container file. It is the first argument to the program. By default it is Corpus.txt.
	buildSpecialCodonArray();	// Build the special codon array. Corpus files are enumerated. Statistics of the corpus are calculated.
	parseQueryFile(argv[2]);	// Parse the query file. It is the second argument to the program. By default it is Query.txt.
	cutilCheckError(cutStopTimer(timer));	// Stop timing.
	printf("done. %u queries. (%.2f ms)\n", query_count, cutGetTimerValue(timer));

	// Open result file and log file for writing headers.
	result_file_handle = fopen(argv[3], "w");
	fprintf(result_file_handle, "Query Index,Match Index,Sequence Index,Ending Position,Sequence Header\n");
	log_file_handle = fopen(argv[4], "w");
	fprintf(log_file_handle, "Query Index,Pattern,Edit Distance,Number of Matches,Querying Time (ms)\n");

	// Use query_index to enumerate pattern[] and edit_distance[].
	for (unsigned int query_index = 0; query_index < query_count; query_index++)
	{
		current_pattern = pattern[query_index];	// Given a pattern, three variables can be derived: m, mask_array, test_bit.
		m = strlen(current_pattern);	// Pattern length.
		k = edit_distance[query_index];	// Edit distance.
		printf("Query %u: searching for %s with edit distance %u ... ", query_index, current_pattern, k);
		match_count = 0;
		cutilCheckError(cutResetTimer(timer));	// Reset timer.
		cutilCheckError(cutStartTimer(timer));	// Start timing.
		if (m <= 32)
		{
			// Derive mask_array and test_bit.
			memset(mask_array_32, 0, sizeof(unsigned int) * CHARACTER_CARDINALITY);
			for (unsigned int i = 0; i < m; i++)	// Derive the mask array of current pattern.
			{
				unsigned int j = (unsigned int)1 << i;
				if (current_pattern[i] == 'N')
				{
					mask_array_32[0] |= j;
					mask_array_32[1] |= j;
					mask_array_32[2] |= j;
					mask_array_32[3] |= j;
				}
				else
				{
					mask_array_32[encode(current_pattern[i])] |= j;
				}
			}
			mask_array_32[0] ^= MAX_UNSIGNED_INT;
			mask_array_32[1] ^= MAX_UNSIGNED_INT;
			mask_array_32[2] ^= MAX_UNSIGNED_INT;
			mask_array_32[3] ^= MAX_UNSIGNED_INT;
			test_bit_32 = (unsigned int)1 << (m - 1);	// The test bit for determining matches of patterns of length 32.

			// About implementation.
			unsigned int sequence_index;	// Used to enumerate [starting sequence, ending sequence).

			// About agrep algorithm.
			unsigned int character_index;	// Used to enumerate each character of a sequence.
			unsigned int mask_word;	// The mask word of a character from mask array.
			unsigned int r[K + 1];	// The most recent columns of K+1 matching tables.
			unsigned int r0;	// The second most recent column of previous matching table.
			unsigned int r1;	// The        most recent column of previous matching table.
			unsigned int r2;	// The second most recent column of current  matching table.
			unsigned int r3;	// The        most recent column of current  matching table. r3 = function(r0, r1, r2, mask_value);
			unsigned int i;	// Used to enumerate K+1 matching tables.

			// Enumerate [starting sequence, ending sequence)
			for (sequence_index = 0; sequence_index < sequence_count; sequence_index++)
			{
				// A new sequence implies a new procedure of agrep algorithm.
				r[0] = MAX_UNSIGNED_INT;
				for (i = 1; i <= k; i++)
					r[i] = r[i - 1] << 1;	// Initialize k+1 matching tables according to agrep algorithm.
				// Enumerate each character of current sequence.
				for (character_index = sequence_cumulative_length[sequence_index]; character_index < sequence_cumulative_length[sequence_index + 1]; character_index++)
				{
					mask_word = mask_array_32[(scodon[character_index >> 4] >> ((character_index & 15) << 1)) & 3];	// Obtain the mask word of current character from mask array.
					r2 = r[0];
					r3 = (r2 << 1) | mask_word;
					r[0] = r3;
					for (i = 1; i <= k; i++)
					{
						r0 = r2;
						r1 = r3;
						r2 = r[i];
						r3 = ((r2 << 1) | mask_word) & ((r0 & r1) << 1) & r0;
						r[i] = r3;
					}
					if (!(r3 & test_bit_32) /* && (match_count < MAX_MATCH_COUNT) */) // A possible match is found.
					{
						if (match_count < MAX_MATCH_COUNT)	// Only the first MAX_MATCH_COUNT matches will be saved.
						{
							match_sequence [match_count] = sequence_index;
							match_character[match_count] = character_index - sequence_cumulative_length[sequence_index];
						}
						match_count++;
					}
				}
			}
		}
		else // m > 32
		{
			// Derive mask_array and test_bit.
			memset(mask_array_64, 0, sizeof(unsigned long long) * CHARACTER_CARDINALITY);
			for (unsigned int i = 0; i < m; i++)	// Derive the mask array of current pattern.
			{
				unsigned long long j = (unsigned long long)1 << i;
				if (current_pattern[i] == 'N')
				{
					mask_array_64[0] |= j;
					mask_array_64[1] |= j;
					mask_array_64[2] |= j;
					mask_array_64[3] |= j;
				}
				else
				{
					mask_array_64[encode(current_pattern[i])] |= j;
				}
			}
			mask_array_64[0] ^= MAX_UNSIGNED_LONG_LONG;
			mask_array_64[1] ^= MAX_UNSIGNED_LONG_LONG;
			mask_array_64[2] ^= MAX_UNSIGNED_LONG_LONG;
			mask_array_64[3] ^= MAX_UNSIGNED_LONG_LONG;
			test_bit_64 = (unsigned long long)1 << (m - 1);	// The test bit for determining matches of patterns of length 64.

			// About implementation.
			unsigned int sequence_index;	// Used to enumerate [starting sequence, ending sequence).

			// About agrep algorithm.
			unsigned int character_index;	// Used to enumerate each character of a sequence.
			unsigned long long mask_word;	// The mask word of a character from mask array.
			unsigned long long r[K + 1];	// The most recent columns of K+1 matching tables.
			unsigned long long r0;	// The second most recent column of previous matching table.
			unsigned long long r1;	// The        most recent column of previous matching table.
			unsigned long long r2;	// The second most recent column of current  matching table.
			unsigned long long r3;	// The        most recent column of current  matching table. r3 = function(r0, r1, r2, mask_value);
			unsigned int i;	// Used to enumerate K+1 matching tables.

			// Enumerate [starting sequence, ending sequence)
			for (sequence_index = 0; sequence_index < sequence_count; sequence_index++)
			{
				// A new sequence implies a new procedure of agrep algorithm.
				r[0] = MAX_UNSIGNED_LONG_LONG;
				for (i = 1; i <= k; i++)
					r[i] = r[i - 1] << 1;	// Initialize k+1 matching tables according to agrep algorithm.
				// Enumerate each character of current sequence.
				for (character_index = sequence_cumulative_length[sequence_index]; character_index < sequence_cumulative_length[sequence_index + 1]; character_index++)
				{
					mask_word = mask_array_64[(scodon[character_index >> 4] >> ((character_index & 15) << 1)) & 3];	// Obtain the mask word of current character from mask array.
					r2 = r[0];
					r3 = (r2 << 1) | mask_word;
					r[0] = r3;
					for (i = 1; i <= k; i++)
					{
						r0 = r2;
						r1 = r3;
						r2 = r[i];
						r3 = ((r2 << 1) | mask_word) & ((r0 & r1) << 1) & r0;
						r[i] = r3;
					}
					if (!(r3 & test_bit_64) /* && (match_count < MAX_MATCH_COUNT) */) // A possible match is found.
					{
						if (match_count < MAX_MATCH_COUNT)	// Only the first MAX_MATCH_COUNT matches will be saved.
						{
							match_sequence [match_count] = sequence_index;
							match_character[match_count] = character_index - sequence_cumulative_length[sequence_index];
						}
						match_count++;
					}
				}
			}
		}
		cutilCheckError(cutStopTimer(timer));	// Stop timing.
		query_time = cutGetTimerValue(timer);	// Obtain the querying time of current query.
		printf("done. %u matches found. (%.2f ms)\n", match_count, query_time);
		fprintf(log_file_handle, "%u,%s,%u,%u,%.2f\n", query_index, current_pattern, k, match_count, query_time);
		if (match_count > MAX_MATCH_COUNT) match_count = MAX_MATCH_COUNT;	// If the number of matches exceeds MAX_MATCH_COUNT, only the first MAX_MATCH_COUNT matches will be saved into the result file.
		for (unsigned int i = 0; i < match_count; i++)
			fprintf(result_file_handle, "%u,%u,%u,%u,%s\n", query_index, i, match_sequence[i], match_character[i], sequence_header[match_sequence[i]]);
	}
	// All queries are done. Release resources.
	fclose(log_file_handle);
	fclose(result_file_handle);
	cutilCheckError(cutDeleteTimer(timer));
	return 0;
}